import com.google.gson.Gson;
import com.google.gson.JsonObject;
import io.milvus.param.MetricType;
import io.milvus.v2.client.MilvusClientV2;
import io.milvus.v2.common.IndexParam;
import io.milvus.v2.service.vector.request.*;
import io.milvus.v2.service.vector.request.data.BaseVector;
import io.milvus.v2.service.vector.request.data.EmbeddedText;
import io.milvus.v2.service.vector.request.data.FloatVec;
import io.milvus.v2.service.vector.request.ranker.BaseRanker;
import io.milvus.v2.service.vector.request.ranker.RRFRanker;
import io.milvus.v2.service.vector.response.DeleteResp;
import io.milvus.v2.service.vector.response.InsertResp;
import io.milvus.v2.service.vector.response.SearchResp;
import lombok.Builder;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.BatchingStrategy;
import org.springframework.ai.embedding.EmbeddingOptionsBuilder;
import org.springframework.ai.observation.conventions.VectorStoreSimilarityMetric;
import org.springframework.ai.openai.OpenAiEmbeddingModel;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.ai.vectorstore.filter.Filter;
import org.springframework.ai.vectorstore.filter.FilterExpressionConverter;
import org.springframework.ai.vectorstore.milvus.MilvusFilterExpressionConverter;
import org.springframework.stereotype.Component;
import org.springframework.util.Assert;

import java.util.*;

/**
 * @description: 基于 milvusV2 的 VectorStore
 * @author: wxw
 * @date: 2025/9/8
 */
@Builder
@Component("milvusVectorStoreV2")
@Slf4j
public class MilvusVectorStoreV2 implements VectorStore {
    private final MilvusClientV2 milvusClient;
    private final OpenAiEmbeddingModel embeddingModel;
    protected final BatchingStrategy batchingStrategy;
    public final FilterExpressionConverter filterExpressionconverter = new MilvusFilterExpressionConverter();
    private static final Map<MetricType, VectorStoreSimilarityMetric> SIMILARITY_TYPE_MAPPING = null;
    private final MetricType metricType = MetricType.COSINE;
    private static final Gson gson = new Gson();
    public final FilterExpressionConverter filterExpressionConverter = new MilvusFilterExpressionConverter();

    /**
     * 密集向量 collection 名称
     */
    private String dense_collectionName = "text_dense_collection";
    /**
     * 稀疏向量 collection 名称
     */
    private String sparse_collectionName = "text_sparse_collection";
    /**
     * 多向量 collection 名称
     */
    private String multi_collectionName = "multi_collection_01";

    /**
     * 密集向量名称
     */
    private String vectorFieldName = "text_dense";

    /**
     * 稀疏向量名称
     */
    private String sparseVectorFieldName = "text_sparse";
    /**
     * 分区名称
     */
    private String partitionName = "_default";


    public MilvusVectorStoreV2(MilvusClientV2 milvusClient,
                               OpenAiEmbeddingModel embeddingModel,
                               BatchingStrategy batchingStrategy) {
        this.milvusClient = milvusClient;
        this.embeddingModel = embeddingModel;
        this.batchingStrategy = batchingStrategy;
    }

    public MilvusVectorStoreV2(MilvusClientV2 milvusClient,
                               OpenAiEmbeddingModel embeddingModel,
                               BatchingStrategy batchingStrategy,
                               String dense_collectionName,
                               String sparse_collectionName,
                               String multi_collectionName,
                               String vectorFieldName,
                               String sparseVectorFieldName,
                               String partitionName) {
        this.milvusClient = milvusClient;
        this.embeddingModel = embeddingModel;
        this.batchingStrategy = batchingStrategy;
        this.dense_collectionName = dense_collectionName;
        this.sparse_collectionName = sparse_collectionName;
        this.multi_collectionName = multi_collectionName;
        this.vectorFieldName = vectorFieldName;
        this.sparseVectorFieldName = sparseVectorFieldName;
        this.partitionName = partitionName;
    }

    @Override
    public String getName() {
        return VectorStore.super.getName();
    }

    /**
     * 插入密集向量到 collection
     * @param documents
     */
    @Override
    public void add(List<Document> documents) {
        Assert.notNull(documents, "Documents must not be null");

        List<JsonObject> data = new ArrayList();
        Gson gson = new Gson();
        int index = 0;

        List<float[]> embeddings = this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(), this.batchingStrategy);

        for(Document document : documents) {
            JsonObject docJson = new JsonObject();

            // id 会自动生成, 不用指定

            // 添加文本内容 text
            docJson.addProperty("text", document.getText());

            // 添加元数据 metadata
            Map<String, Object> metadata = document.getMetadata();
            if (metadata != null) {
                docJson.add("metadata", gson.toJsonTree(metadata));
            } else {
                docJson.add("metadata", new JsonObject());
            }

            // 添加密集向量 text_dense
            float[] embedding = embeddings.get(index);
            List<Float> vectorList = new ArrayList<>();
            for (float value : embedding) {
                vectorList.add(value);
            }
            docJson.add("text_dense", gson.toJsonTree(vectorList));

            data.add(docJson);
            index++;
        }

        InsertReq insertReq = InsertReq.builder()
                .data(data)
                .collectionName(dense_collectionName)
                .partitionName(partitionName)
                .build();

        InsertResp insertResp = this.milvusClient.insert(insertReq);
        if (insertResp.getInsertCnt() == 0) {
            throw new RuntimeException("Failed to insert");
        }
    }

    @Override
    public void accept(List<Document> documents) {
        this.add(documents);
    }

    /**
     * 插入多向量文档到 Milvus
     * @param documents
     */
    public void add_multi(List<Document> documents, String collectionName) {
        Assert.notNull(documents, "Documents must not be null");

        List<JsonObject> data = new ArrayList();
        Gson gson = new Gson();
        int index = 0;

        List<float[]> embeddings = this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(), this.batchingStrategy);

        for(Document document : documents) {
            JsonObject docJson = new JsonObject();

            // id 和 text_parse 会自动生成

            // 添加文本内容 text
            docJson.addProperty("text", document.getText());

            // 添加密集向量 text_dense
            float[] embedding = embeddings.get(index);
            List<Float> vectorList = new ArrayList<>();
            for (float value : embedding) {
                vectorList.add(value);
            }
            docJson.add("text_dense", gson.toJsonTree(vectorList));

            data.add(docJson);
            index++;
        }

        InsertReq insertReq = InsertReq.builder()
                .data(data)
                .collectionName(collectionName)
                .partitionName(this.partitionName)
                .build();

        InsertResp insertResp = this.milvusClient.insert(insertReq);
        if (insertResp.getInsertCnt() == 0) {
            throw new RuntimeException("Failed to insert");
        }
    }


    /**
     * 从密集向量数据库删除文档(无 filterExpression)
     * @param idList
     */
    @Override
    public void delete(List<String> idList) {
        Assert.notNull(idList, "IdList must not be null");

        // 2. 将 List<String> 安全转为 List<Long>
        List<Long> longIdList = new ArrayList<>(idList.size());
        for (String strId : idList) {
            // 校验单个 ID
            if (strId == null || strId.trim().isEmpty()) {
                throw new IllegalArgumentException("Invalid ID: empty or null string");
            }
            // 字符串转 Long
            try {
                Long longId = Long.parseLong(strId.trim()); // 去除空格后转换
                longIdList.add(longId);
            } catch (NumberFormatException e) {
                throw new IllegalArgumentException("Failed to convert ID to Long: " + strId, e);
            }
        }

        DeleteReq deleteReq = DeleteReq.builder()
                .collectionName(dense_collectionName)
                .partitionName("_default")
                .ids(new ArrayList<>(longIdList))
                .build();
        DeleteResp delete = this.milvusClient.delete(deleteReq);

        if (delete.getDeleteCnt() == 0) {
            throw new RuntimeException("Failed to delete");
        }
    }

    /**
     * 从密集向量数据库删除文档(使用 filterExpression)
     * @param filterExpression
     */
    @Override
    public void delete(Filter.Expression filterExpression) {
        Assert.notNull(filterExpression, "Filter expression must not be null");

        try {
            String nativeFilterExpression = this.filterExpressionConverter.convertExpression(filterExpression);
            DeleteResp deleteResp = this.milvusClient.delete(
                    DeleteReq.builder()
                            .collectionName(dense_collectionName)
                            .partitionName("_default")
                            .filterTemplateValues(Map.of("filter", nativeFilterExpression))
                            .build()
            );

            if (deleteResp.getDeleteCnt() == 0) {
                throw new IllegalStateException("Failed to delete documents by filter");
            } else {
                long deleteCount = deleteResp.getDeleteCnt();
                log.debug("Deleted {} documents matching filter expression", deleteCount);
            }
        } catch (Exception e) {
            log.error("Failed to delete documents by filter: {}", e.getMessage(), e);
            throw new IllegalStateException("Failed to delete documents by filter", e);
        }
    }

    /**
     * 从密集向量数据库删除文档(使用 filterExpression-String 类型)
     * @param filterExpression
     */
    @Override
    public void delete(String filterExpression) {
        Assert.notNull(filterExpression, "Filter expression must not be null");

        try {
            DeleteResp deleteResp = this.milvusClient.delete(
                    DeleteReq.builder()
                            .collectionName(dense_collectionName)
                            .partitionName("_default")
                            .filterTemplateValues(Map.of("filter", filterExpression))
                            .build()
            );

            if (deleteResp.getDeleteCnt() == 0) {
                throw new IllegalStateException("Failed to delete documents by filter");
            } else {
                long deleteCount = deleteResp.getDeleteCnt();
                log.debug("Deleted {} documents matching filter expression", deleteCount);
            }
        } catch (Exception e) {
            log.error("Failed to delete documents by filter: {}", e.getMessage(), e);
            throw new IllegalStateException("Failed to delete documents by filter", e);
        }
    }

    /**
     * 语义检索-密集向量
     * @param request
     * @return
     */
    @Override
    public List<Document> similaritySearch(SearchRequest request) {
        if (request == null) {
            throw new NullPointerException("requestParam is marked non-null but is null");
        }

        // 提取参数
        float[] userQueryEmbedding = this.getUserQueryEmbedding(request.getQuery());
        int topK = request.getTopK();
        List<String> outputFields = Arrays.asList("id","content","metadata");
        // 构造过滤条件
        Filter.Expression filterExpression = request.getFilterExpression();
        String nativeFilterExpression = this.filterExpressionConverter.convertExpression(filterExpression);

        // 构造问题向量
        List<BaseVector> vectors = new ArrayList<>();
        FloatVec floatVec = new FloatVec(userQueryEmbedding);
        vectors.add(floatVec);

        // 构造 SearchReq
        SearchReq searchReq = SearchReq.builder()
                .collectionName(dense_collectionName)
                .partitionNames(Collections.singletonList(this.partitionName))
                .metricType(IndexParam.MetricType.COSINE)
                .annsField(this.vectorFieldName)
                .topK(topK)
                .filter(nativeFilterExpression)
                .outputFields(outputFields)
                .data(vectors)
                .build();

        // 执行检索
        SearchResp searchResp = milvusClient.search(searchReq);

        List<Document> documents = this.searchRespToDocument(searchResp);
        return documents;
    }


    /**
     * 语义检索-密集向量
     * @param query
     * @return
     */
    @Override
    public List<Document> similaritySearch(String query) {
        if (query == null) {
            throw new NullPointerException("query is marked non-null but is null");
        }

        float[] userQueryEmbedding = this.getUserQueryEmbedding(query);
        final int topK = 5;
        List<String> outputFields = Arrays.asList("id","content","metadata");

        // 构造问题向量
        List<BaseVector> vectors = new ArrayList<>();
        FloatVec floatVec = new FloatVec(userQueryEmbedding);
        vectors.add(floatVec);

        // 构造向量检索请求
        SearchReq searchReq = SearchReq.builder()
                .collectionName(dense_collectionName)
                .partitionNames(Collections.singletonList(this.partitionName))
                .metricType(IndexParam.MetricType.COSINE)
                .annsField(this.vectorFieldName)
                .topK(topK)
                .outputFields(outputFields)
                .data(vectors)
                .build();

        // 执行检索
        SearchResp searchResp = milvusClient.search(searchReq);

        List<Document> documents = this.searchRespToDocument(searchResp);
        return documents;
    }

    /**
     * 使用多路召回(单向量库)+RFF的向量检索
     * @param request
     * @return
     */
    public List<Document> multiChannelRecallSearch(SearchRequest request, String keywords) {
        if (request == null) {
            throw new NullPointerException("requestParam is marked non-null but is null");
        }

        // 提取参数
        float[] userQueryEmbedding = this.getUserQueryEmbedding(request.getQuery());
        int topK = request.getTopK();
        Filter.Expression filterExpression = request.getFilterExpression();
        String nativeFilterExpression = this.filterExpressionConverter.convertExpression(filterExpression);

        List<String> outputFields = Arrays.asList("id","text");

        // 构造用户输入文本向量
        List<BaseVector> queryDenseVectors = new ArrayList<>();
        FloatVec floatVec = new FloatVec(userQueryEmbedding);
        queryDenseVectors.add(floatVec);

        // 构造全文检索参数
        List<BaseVector> queryTexts = Collections.singletonList(new EmbeddedText(keywords));

        // 构造多路召回检索请求
        List<AnnSearchReq> searchRequests = new ArrayList<>();
        if(nativeFilterExpression.isEmpty()) {
            searchRequests.add(AnnSearchReq.builder()
                    .vectorFieldName("text_dense")
                    .vectors(queryDenseVectors)
                    .params("{\"nprobe\": 10}") // nprobe 参数决定了在查询过程中需要检查的倒排列表的数量
                    .topK(5)
                    .build());
        } else {
            searchRequests.add(AnnSearchReq.builder()
                    .vectorFieldName("text_dense")
                    .vectors(queryDenseVectors)
                    .params("{\"nprobe\": 10}") // nprobe 参数决定了在查询过程中需要检查的倒排列表的数量
                    .expr(nativeFilterExpression) // 筛选条件
                    .topK(5)
                    .build());
        }
        searchRequests.add(AnnSearchReq.builder()
                .vectorFieldName("text_sparse")
                .vectors(queryTexts)
                .params("{\"drop_ratio_search\": 0.2}") // 控制搜索时过滤掉的稀疏向量维度比例
                .topK(3)
                .build());

        // 互易排名融合排名器 RRFRanker
        BaseRanker reranker = new RRFRanker(100);

        // 构造混合检索请求
        HybridSearchReq hybridSearchReq = HybridSearchReq.builder()
                .collectionName(multi_collectionName)
                .partitionNames(Collections.singletonList(partitionName))
                .searchRequests(searchRequests)
                .ranker(reranker)
                .outFields(outputFields)
                .topK(topK)
                .build();

        // 执行混合检索
        SearchResp searchResp = milvusClient.hybridSearch(hybridSearchReq);

        List<Document> documents = this.searchRespToDocument(searchResp);
        return documents;
    }

    /**
     * 使用多路召回(多向量库)+RRF+Reranker的向量检索
     * @param request
     * @return
     */
    public List<Document> multiChannelRecallAndDataSourceSearch(SearchRequest request) {
        // TODO
        return null;
    }

    /**
     * 全文检索(关键词检索)
     * @param request
     * @return
     */
    public List<Document> fullTextSearch(SearchRequest request) {

        int topK = request.getTopK();
        String query = request.getQuery();
        List<String> outputFields = Arrays.asList("id","text");

        // 构造全文检索参数
        List<BaseVector> data = Collections.singletonList(new EmbeddedText(query));

        // 全文搜索参数
        Map<String,Object> searchParams = new HashMap<>();
        searchParams.put("drop_ratio_search", 0.2);

        // 执行全文搜索
        SearchResp searchResp = milvusClient.search(SearchReq.builder()
                .collectionName(sparse_collectionName)
                .data(data)
                .annsField(sparseVectorFieldName)
                .topK(topK)
                .searchParams(searchParams)
                .outputFields(outputFields)
                .build());

        List<Document> documents = this.searchRespToDocument(searchResp);
        return documents;
    }

    /**
     * 获取用户查询的向量嵌入
     * @param query
     * @return
     */
    private float[] getUserQueryEmbedding(String query) {
        return this.embeddingModel.embed(query);
    }

    
    public void setDense_collectionName(String dense_collectionName) {
        this.dense_collectionName = dense_collectionName;
    }

    public void setSparse_collectionName(String sparse_collectionName) {
        this.sparse_collectionName = sparse_collectionName;
    }

    public void setMulti_collectionName(String multi_collectionName) {
        this.multi_collectionName = multi_collectionName;
    }

    public void setVectorFieldName(String vectorFieldName) {
        this.vectorFieldName = vectorFieldName;
    }

    public void setSparseVectorFieldName(String sparseVectorFieldName) {
        this.sparseVectorFieldName = sparseVectorFieldName;
    }

    public void setPartitionName(String partitionName) {
        this.partitionName = partitionName;
    }

    /**
     * 转换 SearchResp 为 Document
     * @param searchResp
     * @return
     */
    private List<Document> searchRespToDocument(SearchResp searchResp) {
        List<Document> documents = new ArrayList<>();
        List<List<SearchResp.SearchResult>> searchResults = searchResp.getSearchResults();
        for (List<SearchResp.SearchResult> results : searchResults) {
            for (SearchResp.SearchResult result : results) {
                Map<String, Object> entity = result.getEntity();
                Document document = Document.builder()
                        .id(String.valueOf(result.getId()))
                        .score(Double.valueOf(result.getScore()))
                        .metadata(entity)
                        .text(entity.get("text").toString())
                        .build();
                documents.add(document);
                log.info("Document: {}", document);
            }
        }
        return documents;
    }
}

